#include <windows.h>
#include <stdio.h>
#pragma warning(disable:4189)	//we want GetLastError() stored
								//so we can see it in a crash dump
								//but we don't use it otherwise.
#include "lib.h"
#include "protectedmemory.h"

#if defined(_INC_CRTDBG) && defined(_CRTDBG_MAP_ALLOC)
#undef malloc
#undef free
#undef new
#endif

#ifdef PM_LOCAL_PROTECTION

#define PM_PAGESIZE				4096
#define PM_FINGERPRINT			((uint64)0xaaaaaaaaaaaaaaaa)
#define PM_FOOTPRINT			((uint64)0xbbbbbbbbbbbbbbbb)
#define PM_FLAGS_ALLOCATED		((uint64)0x0000000000000000)
#define PM_FLAGS_DEALLOCATED	((uint64)0xffffffffffffffff)


struct ProtectedMemory::Head
{
	uint64		fingerPrint;
	char *		baseAddress;
	Head *		headAddress;
	char *		userAddress;
	char *		lockAddress;
	uint64		flags;
	uint64		footPrint;
};

void *
ProtectedMemory::protected_malloc(size_t size)
{
	size_t	pageSize;
	Head	head;
	uint32	crap;

	if (size < 0)
		size = 0;

	pageSize  = sizeof(Head);				// The head data.
	pageSize += size;						// The user data.
	pageSize += PM_PAGESIZE;				// The lock data.

	// Now round up to the next page boundary.
	pageSize  = (pageSize + PM_PAGESIZE - 1) & ~(PM_PAGESIZE - 1);

	// Generate our head.
	head.fingerPrint = PM_FINGERPRINT;
	head.baseAddress = (char *)VirtualAlloc(NULL, pageSize, MEM_COMMIT, PAGE_READWRITE);
	head.lockAddress = head.baseAddress + pageSize - PM_PAGESIZE;
	head.userAddress = head.lockAddress - size;
	head.headAddress = (Head *)(head.userAddress - sizeof(Head));
	head.flags       = PM_FLAGS_ALLOCATED;
	head.footPrint   = PM_FOOTPRINT;

	// Protect the trailing 'lock' page (protection against buffer overruns).
	if (head.baseAddress)
	{
		if (VirtualProtect(head.lockAddress, PM_PAGESIZE, PAGE_NOACCESS, &crap))
		{
			// Copy the head data into the allocated memory.
			*head.headAddress = head;

			// Return a usable pointer.
			return head.userAddress;
		}
		else
		{
			// A serious system-level error has occurred.  Grab the error
			// code and halt execution for postmortem.
			volatile int32 gle = GetLastError();

			__asm { int 3 }

			return 0;
		}
	}
	else
	{
		// A serious system-level error has occurred.  Grab the error
		// code and halt execution for postmortem.
		volatile int32 gle = GetLastError();

		__asm { int 3 }

		return 0;
	}
}


void
ProtectedMemory::protected_free(void * memory)
{
	Head *	head;

	// Find our head.
	head = (Head *)((char *)memory - sizeof(Head));

	// Make sure it's our head (and release).
	if (memory &&
		head->fingerPrint == PM_FINGERPRINT &&
		head->footPrint   == PM_FOOTPRINT &&
		head->headAddress == head &&
		head->userAddress == memory &&
		head->flags       == PM_FLAGS_ALLOCATED)
	{
		head->flags = PM_FLAGS_DEALLOCATED;
		if (!VirtualFree(head->baseAddress, 0, MEM_RELEASE))
		{
			// A serious system-level error has occurred.  Grab the error
			// code and halt execution for postmortem.
			volatile int32 gle = GetLastError();

			__asm { int 3 }
		}
	}
	else
	{
		// A serious developer-level error has occurred.  Halt execution
		// for postmortem.
		//
		// Possible reasons:
		//    1. Null-pointer deallocation (memory == 0)
		//    2. Double deallocation.      (head->flags == PM_FLAGS_DEALLOCATED)
		//    3. Buffer underrun.          (head->footPrint != PM_FOOTPRINT)
		//    5. Invalid pointer.          (head->userAddress != memory)
		__asm { int 3 }
	}
}
#endif	// PM_LOCAL_PROTECTION


//----------------------------------------------------------------------------
// Global protection.
// Warning: Use these global routines ONLY on systems with LOTS of memory!
//
#ifdef PM_GLOBAL_PROTECTION
void * __cdecl
operator new(size_t size) _THROW1(std::bad_alloc)
{
	void * ptr = ProtectedMemory::protected_malloc(size);

	if (!ptr)
		throw(std::bad_alloc());

	return ptr;
}

void * __cdecl
operator new(size_t size, const std::nothrow_t &)
{
	return ProtectedMemory::protected_malloc(size);
}

void * __cdecl
operator new[](size_t size) _THROW1(std::bad_alloc)
{
	void * ptr = ProtectedMemory::protected_malloc(size);

	if (!ptr)
		throw(std::bad_alloc());

	return ptr;
}

void * __cdecl
operator new[](size_t size, const std::nothrow_t &)
{
	return ProtectedMemory::protected_malloc(size);
}

void __cdecl
operator delete(void * ptr) _THROW0()
{
	// Deleting a null-pointer is *not* an error.  Do not pass the pointer.
	if (ptr == 0)
		return;

	ProtectedMemory::protected_free(ptr);
}

void __cdecl
operator delete[](void * ptr) _THROW0()
{
	// Deleting an array null-pointer is *not* an error.  Do not pass the pointer.
	if (ptr == 0)
		return;

	ProtectedMemory::protected_free(ptr);
}

void *
malloc(size_t size)
{
	return ProtectedMemory::protected_malloc(size);
}

void
free(void * ptr)
{
	// Null-pointer free() errors are handled by ProtectedMemory::protected_free().
	ProtectedMemory::protected_free(ptr);
}
#endif	// PM_GLOBAL_PROTECTION


//
// vim: nowrap
//
